# -*- coding: utf-8 -*-

import torch
import random
import numpy as np
# import segmentation_models_pytorch as smp
from transform import random_transform


class TCSM(torch.nn.Module):
    def __init__(self, student_model, teacher_model, student_device='cuda:0', teacher_device='cuda:0', smoothing=0.999):
        super(TCSM, self).__init__()
        
        self.student_model = student_model
        self.teacher_model = teacher_model
        for param in self.teacher_model.parameters():
            param.detach_()
        
        self.to_device(student_device, teacher_device)
        
        self.smoothing = smoothing
        np.random.seed(random.randint(1, 10000))
    
    def to_device(self, student_device, teacher_device):
        self.student_model.to(student_device)
        self.teacher_model.to(teacher_device)
        
        self.student_device = student_device
        self.teacher_device = teacher_device
        
    def forward(self, x):
        x = x.float()
        
        x = x.to(self.student_device)
        student_output = self.student_model(x)
        
        x = x.to(self.teacher_device)
        x, trans_args = random_transform(x)
        teacher_output = self.teacher_model(x)
        
        return student_output, teacher_output, trans_args
    
    def predict(self, x):
        x = x.float()
        x = x.to(self.student_device)
        
        with torch.no_grad():
            pred = self.student_model(x)
        
        return pred
        

# if __name__ == "__main__":
#     from CC_CCI_dataset import LabeledDataset, UnlabeledDataset
#     # import matplotlib.pyplot as plt
#     IMAGE_DIR = "D:\data\ct_lesion_seg\image"
#     MASK_DIR = "D:\data\ct_lesion_seg\mask"
#     labeled_dataset = LabeledDataset(IMAGE_DIR, MASK_DIR)
#     unlabeled_dataset = UnlabeledDataset(IMAGE_DIR, MASK_DIR, len(labeled_dataset))
#     labeled_loader = torch.utils.data.DataLoader(labeled_dataset, batch_size=8, shuffle=True)
#     unlabeled_loader = torch.utils.data.DataLoader(unlabeled_dataset, batch_size=8, shuffle=True)
    
#     student = smp.DeepLabV3Plus(encoder_weights=None, in_channels=1, classes=3, activation='softmax2d')
#     teacher = smp.DeepLabV3Plus(encoder_weights=None, in_channels=1, classes=3, activation='softmax2d')
#     tcsm = TCSM(student, teacher)
    
#     labeled_loader_iter = iter(labeled_loader)
#     unlabeled_loader_iter = iter(unlabeled_loader)
#     for i in range(1):
#         x, y = labeled_loader_iter.next()
#         z = unlabeled_loader_iter.next()
        
#         ls, lt, args = tcsm(x)
#         # student(x.cuda())
#         us, ut, args = tcsm(y)